from trdg.data_generator import FakeTextDataGenerator
import random
# import sys
import os
# sys.path.append(r"/home/gsh/PycharmProjects/TextRecognitionDataGenerator/trdg")
from trdg.utils import load_fonts
from tqdm import tqdm


class TRDG_TOOL():
    def __init__(self, blur_range, txt_path, text_length_range, colors, skewing_angle, count, font=[], size=32,
                 language='77'):
        # colors format -> ['red', 'black', ...]
        # skewing->[0-skewing,skewing]

        self.fonts = self.fonts_filter(language, font)
        self.blur_range = self.range_filter(blur_range)
        self.init_text_from_file(txt_path)  # 把一个方法钓上来，穿个参数进去
        self.text_length_range = self.range_filter(text_length_range)
        self.size = size
        # self.color_lookup_dict = self.gen_color_lookup_dict()
        # self.color_options = self.color_filter(colors)
        self.swiking_angle = self.skewing_filter(skewing_angle)
        self.background_type = 3  # Set background with picture.
        self.generated_count = 0
        self.count = self.count_filter(count)
        ###############################################
        self.distorsion_type = 0
        self.distorsion_orientation = 0
        self.is_handwritten = False
        self.width = -1
        self.alignment = 1
        self.orientation = 0
        self.space_width = 1.0
        self.character_spacing = 0
        self.margins = (5, 5, 5, 5)
        self.fit = False
        self.output_mask = False
        self.word_split = False

        self.colormap = {
            # X11 colour table from https://drafts.csswg.org/css-color-4/, with
            # gray/grey spelling issues fixed.  This is a superset of HTML 4.0
            # colour names used in CSS 1.
            "aliceblue": "#f0f8ff",
            "antiquewhite": "#faebd7",
            "aqua": "#00ffff",
            "aquamarine": "#7fffd4",
            "azure": "#f0ffff",
            "beige": "#f5f5dc",
            "bisque": "#ffe4c4",
            "black": "#000000",
            "blanchedalmond": "#ffebcd",
            "blue": "#0000ff",
            "blueviolet": "#8a2be2",
            "brown": "#a52a2a",
            "burlywood": "#deb887",
            "cadetblue": "#5f9ea0",
            "chartreuse": "#7fff00",
            "chocolate": "#d2691e",
            "coral": "#ff7f50",
            "cornflowerblue": "#6495ed",
            "cornsilk": "#fff8dc",
            "crimson": "#dc143c",
            "cyan": "#00ffff",
            "darkblue": "#00008b",
            "darkcyan": "#008b8b",
            "darkgoldenrod": "#b8860b",
            "darkgray": "#a9a9a9",
            "darkgrey": "#a9a9a9",
            "darkgreen": "#006400",
            "darkkhaki": "#bdb76b",
            "darkmagenta": "#8b008b",
            "darkolivegreen": "#556b2f",
            "darkorange": "#ff8c00",
            "darkorchid": "#9932cc",
            "darkred": "#8b0000",
            "darksalmon": "#e9967a",
            "darkseagreen": "#8fbc8f",
            "darkslateblue": "#483d8b",
            "darkslategray": "#2f4f4f",
            "darkslategrey": "#2f4f4f",
            "darkturquoise": "#00ced1",
            "darkviolet": "#9400d3",
            "deeppink": "#ff1493",
            "deepskyblue": "#00bfff",
            "dimgray": "#696969",
            "dimgrey": "#696969",
            "dodgerblue": "#1e90ff",
            "firebrick": "#b22222",
            "floralwhite": "#fffaf0",
            "forestgreen": "#228b22",
            "fuchsia": "#ff00ff",
            "gainsboro": "#dcdcdc",
            "ghostwhite": "#f8f8ff",
            "gold": "#ffd700",
            "goldenrod": "#daa520",
            "gray": "#808080",
            "grey": "#808080",
            "green": "#008000",
            "greenyellow": "#adff2f",
            "honeydew": "#f0fff0",
            "hotpink": "#ff69b4",
            "indianred": "#cd5c5c",
            "indigo": "#4b0082",
            "ivory": "#fffff0",
            "khaki": "#f0e68c",
            "lavender": "#e6e6fa",
            "lavenderblush": "#fff0f5",
            "lawngreen": "#7cfc00",
            "lemonchiffon": "#fffacd",
            "lightblue": "#add8e6",
            "lightcoral": "#f08080",
            "lightcyan": "#e0ffff",
            "lightgoldenrodyellow": "#fafad2",
            "lightgreen": "#90ee90",
            "lightgray": "#d3d3d3",
            "lightgrey": "#d3d3d3",
            "lightpink": "#ffb6c1",
            "lightsalmon": "#ffa07a",
            "lightseagreen": "#20b2aa",
            "lightskyblue": "#87cefa",
            "lightslategray": "#778899",
            "lightslategrey": "#778899",
            "lightsteelblue": "#b0c4de",
            "lightyellow": "#ffffe0",
            "lime": "#00ff00",
            "limegreen": "#32cd32",
            "linen": "#faf0e6",
            "magenta": "#ff00ff",
            "maroon": "#800000",
            "mediumaquamarine": "#66cdaa",
            "mediumblue": "#0000cd",
            "mediumorchid": "#ba55d3",
            "mediumpurple": "#9370db",
            "mediumseagreen": "#3cb371",
            "mediumslateblue": "#7b68ee",
            "mediumspringgreen": "#00fa9a",
            "mediumturquoise": "#48d1cc",
            "mediumvioletred": "#c71585",
            "midnightblue": "#191970",
            "mintcream": "#f5fffa",
            "mistyrose": "#ffe4e1",
            "moccasin": "#ffe4b5",
            "navajowhite": "#ffdead",
            "navy": "#000080",
            "oldlace": "#fdf5e6",
            "olive": "#808000",
            "olivedrab": "#6b8e23",
            "orange": "#ffa500",
            "orangered": "#ff4500",
            "orchid": "#da70d6",
            "palegoldenrod": "#eee8aa",
            "palegreen": "#98fb98",
            "paleturquoise": "#afeeee",
            "palevioletred": "#db7093",
            "papayawhip": "#ffefd5",
            "peachpuff": "#ffdab9",
            "peru": "#cd853f",
            "pink": "#ffc0cb",
            "plum": "#dda0dd",
            "powderblue": "#b0e0e6",
            "purple": "#800080",
            "rebeccapurple": "#663399",
            "red": "#ff0000",
            "rosybrown": "#bc8f8f",
            "royalblue": "#4169e1",
            "saddlebrown": "#8b4513",
            "salmon": "#fa8072",
            "sandybrown": "#f4a460",
            "seagreen": "#2e8b57",
            "seashell": "#fff5ee",
            "sienna": "#a0522d",
            "silver": "#c0c0c0",
            "skyblue": "#87ceeb",
            "slateblue": "#6a5acd",
            "slategray": "#708090",
            "slategrey": "#708090",
            "snow": "#fffafa",
            "springgreen": "#00ff7f",
            "steelblue": "#4682b4",
            "tan": "#d2b48c",
            "teal": "#008080",
            "thistle": "#d8bfd8",
            "tomato": "#ff6347",
            "turquoise": "#40e0d0",
            "violet": "#ee82ee",
            "wheat": "#f5deb3",
            "white": "#ffffff",
            "whitesmoke": "#f5f5f5",
            "yellow": "#ffff00",
            "yellowgreen": "#9acd32",
        }
        self.color = self.color_filter(colors)

    def count_filter(self, my_count):
        if isinstance(my_count, int):
            if my_count < 1:
                return self.err_input()
            else:
                return my_count

    # def gen_color_lookup_dict(self):
    #     # color lookup dict format :
    #     # {color_name : [[start_red, end_red], [start_green, end_green], [start_blue, end_blue]], ...}
    #     color_lookup_dict = {'red': [[240, 255], [0, 20], [0, 20]], 'green': [[0, 20], [240, 255], [0, 20]]}
    #     return color_lookup_dict

    def skewing_filter(self, skewing):
        if isinstance(skewing, int):
            return skewing
        else:
            return self.err_input()

    def color_filter(self, input_colors):

        res_color = []
        if isinstance(input_colors, list):

            for i in input_colors:
                if i in self.colormap:
                    res_color.append(i)

                else:
                    raise ValueError("unknown color specifier: %s" % input_colors)
            return res_color

        if isinstance(input_colors, str):

            if input_colors in self.colormap:
                return list(input_colors)

    def gen_random_color(self, colors):
        colors_name = random.choice(colors)
        return self.colormap.get(colors_name)

    # def color_filter(self, colors):
    #     colors_accept = []
    #     if isinstance(colors, list):
    #         for color in colors:
    #             if color in self.color_lookup_dict:
    #                 colors_accept.append(color)
    #     elif isinstance(colors, str):
    #         if colors in self.color_lookup_dict:
    #             colors_accept.append(colors)
    #     if not colors_accept:
    #         self.err_input()
    #     return colors_accept

    def fonts_filter(self, language, font_list):
        fonts = []
        if language == 'Null':  # Null的话就用我们自己输入的font,"cn"是cn文件夹里的，"其他"是latin文件夹下的
            if isinstance(font_list, list):
                for font in font_list:
                    if os.path.exists(font):
                        fonts.append(font)
            elif isinstance(font_list, str):
                if os.path.exists(font_list):
                    fonts.append(font_list)
        else:
            fonts = load_fonts(language)

        if not fonts:
            self.err_input()

        return fonts

    def range_filter(self, range_data):
        range_res = [0, 0]
        if isinstance(range_data, list):
            if len(range_data) == 2:
                range_res[0] = min(range_data[0], range_data[1])
                range_res[1] = max(range_data[0], range_data[1])
                return range_res
            elif len(range_data) == 1:
                return range_data[0]
            else:
                self.err_input()
        elif isinstance(range_data, int):
            return range_data
        else:
            self.err_input()

    def err_input(self):
        print("Error input parameters.")

    # def get_text_from_file(self, file_path, line_num = -1):
    #     with open(file_path, 'r') as f:

    def init_text_from_file(self, file_path):
        with open(file_path, 'r') as f:
            self.text_list = f.readlines()
        self.text_list = [i.replace('\n', '') for i in self.text_list]
        self.text_num = len(self.text_list)

    def gen_random_text(self):
        if isinstance(self.text_length_range, int):
            text_length = random.randint(1, self.text_length_range)
        else:
            low, high = self.text_length_range
            text_length = random.randint(low, high)

        line_pick_list = []
        text_res_length = 0
        text_res = ''
        while text_res_length < text_length:
            line_index = random.randint(0, self.text_num - 1)
            if line_index in line_pick_list:
                continue
            line_pick_list.append(line_index)
            line_text = self.text_list[line_index]
            line_length = len(line_text)

            if line_length > text_length - text_res_length:
                start_pos = random.randint(
                    0, line_length - text_res_length - 1)
                end_pos = start_pos + text_length - text_res_length
                text_res += line_text[start_pos:end_pos]
                text_res_length = len(text_res)
                break
            else:
                text_res += line_text
                text_res_length += line_length

        return text_res

    def gen_random_blur(self):
        if isinstance(self.blur_range, int):
            blur_para = random.randint(0, self.blur_range)
        else:
            low, high = self.blur_range
            blur_para = random.randint(low, high)
        return blur_para

    # def gen_random_color(self):
    #     color_pick = random.choice(self.color_options)
    #     color_res = '#'
    #     for i in self.color_lookup_dict[color_pick]:
    #         color_strength = random.randint(i[0], i[1])
    #         color_hex = hex(color_strength)
    #         color_res += color_hex[2:]
    #     return color_res

    # def gen_random_skewing(self):
    #     pass

    def gen_random_font(self):
        return random.choice(self.fonts)

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        # 方法返回值做下面类的输入参数
        text = self.gen_random_text()
        blur = self.gen_random_blur()
        color = self.gen_random_color(self.color)
        font = self.gen_random_font()
        random_blur = True
        random_skew = True
        skewing_angle = self.swiking_angle
        #

        if self.generated_count == self.count:
            raise StopIteration
        self.generated_count += 1

        img = FakeTextDataGenerator.generate(
            index=self.generated_count,
            text=text,
            font=font,
            out_dir=None,
            size=self.size,
            extension=None,
            skewing_angle=skewing_angle,
            random_skew=random_skew,
            blur=blur,
            random_blur=random_blur,
            background_type=self.background_type,
            distorsion_type=self.distorsion_type,
            distorsion_orientation=self.distorsion_orientation,
            is_handwritten=self.is_handwritten,
            name_format=0,
            width=self.width,
            alignment=self.alignment,
            text_color="#281827",
            orientation=0,  # self.orientation,
            space_width=self.space_width,
            character_spacing=self.character_spacing,
            margins=self.margins,
            fit=self.fit,
            output_mask=self.output_mask,
            word_split=self.word_split,
        )
        return (img, text)


if __name__ == "__main__":

    # blur_range=[1,5]
    # txt_path=""
    # text_length_range=[3,10]
    # color = "red"
    # skewing_angle=30
    # count=100
    # font = load_fonts("cn")
    # size = 32,
    # language = 'Null'

    generator = TRDG_TOOL(blur_range=[0, 1],
                          txt_path="cn.txt",
                          text_length_range=[3, 10],
                          colors=['aliceblue', 'antiquewhite', 'aqua', 'aquamarine', 'azure', 'beige', 'bisque',
                                  'black', 'blanchedalmond', 'blue', 'blueviolet', 'brown', 'burlywood', 'cadetblue',
                                  'chartreuse',
                                  'chocolate', 'coral', 'cornflowerblue', 'cornsilk', 'crimson', 'cyan', 'darkblue',
                                  'darkcyan', 'darkgoldenrod',
                                  'darkgray', 'darkgrey', 'darkgreen', 'darkkhaki', 'darkmagenta', 'darkolivegreen',
                                  'darkorange', 'darkorchid', 'darkred',
                                  'darksalmon', 'darkseagreen', 'darkslateblue', 'darkslategray', 'darkslategrey',
                                  'darkturquoise', 'darkviolet', 'deeppink',
                                  'deepskyblue', 'dimgray', 'dimgrey', 'dodgerblue', 'firebrick', 'floralwhite',
                                  'forestgreen', 'fuchsia', 'gainsboro', 'ghostwhite',
                                  'gold', 'goldenrod', 'gray', 'grey', 'green', 'greenyellow', 'honeydew', 'hotpink',
                                  'indianred', 'indigo', 'ivory', 'khaki',
                                  'lavender', 'lavenderblush', 'lawngreen', 'lemonchiffon', 'lightblue', 'lightcoral',
                                  'lightcyan', 'lightgoldenrodyellow',
                                  'lightgreen', 'lightgray', 'lightgrey', 'lightpink', 'lightsalmon', 'lightseagreen',
                                  'lightskyblue', 'lightslategray',
                                  'lightslategrey', 'lightsteelblue', 'lightyellow', 'lime', 'limegreen', 'linen',
                                  'magenta', 'maroon', 'mediumaquamarine', 'mediumblue', 'mediumorchid', 'mediumpurple',
                                  'mediumseagreen', 'mediumslateblue', 'mediumspringgreen', 'mediumturquoise',
                                  'mediumvioletred', 'midnightblue', 'mintcream', 'mistyrose', 'moccasin',
                                  'navajowhite', 'navy', 'oldlace', 'olive', 'olivedrab', 'orange', 'orangered',
                                  'orchid', 'palegoldenrod', 'palegreen', 'paleturquoise', 'palevioletred',
                                  'papayawhip', 'peachpuff', 'peru', 'pink', 'plum', 'powderblue', 'purple',
                                  'rebeccapurple', 'red', 'rosybrown', 'royalblue', 'saddlebrown', 'salmon',
                                  'sandybrown', 'seagreen', 'seashell', 'sienna', 'silver', 'skyblue', 'slateblue',
                                  'slategray', 'slategrey', 'snow', 'springgreen', 'steelblue', 'tan', 'teal',
                                  'thistle',
                                  'tomato', 'turquoise', 'violet', 'wheat', 'white', 'whitesmoke', 'yellow',
                                  'yellowgreen'],
                          skewing_angle=5,
                          count=300,
                          font=load_fonts("cn"),
                          size=60,
                          language='cn')
    total_generate_num = generator.count

    
    with tqdm(total=total_generate_num) as tbar:
        i = 1
        for img, lbl in generator:
            img.save("/home/gsh/PycharmProjects/TextRecognitionDataGenerator/pictures_own/{}.jpg".format(i)) #图片生成路径
            with open("/home/gsh/PycharmProjects/TextRecognitionDataGenerator/pictures_own/{}.txt".format(i), 'w') as f:#标签保存路径
                f.write(lbl)
            tbar.update(1)
            i += 1
            if i > total_generate_num:
                break
